import argparse
import grpc
import numpy as np
from google.protobuf import empty_pb2
import play_with_pybullet.grpc_env.arm_env_pb2  as pb
import play_with_pybullet.grpc_env.arm_env_pb2_grpc as rpc
from  play_with_pybullet.grpc_env.png2mp4 import images_to_video
from time import sleep

def main(args):
    channel = grpc.insecure_channel(f"{args.host}:{args.port}")
    stub = rpc.ArmEnvStub(channel)

    # Query spaces
    spaces = stub.GetSpaces(empty_pb2.Empty())
    print("Action dim:", spaces.action_dim, "low:", list(spaces.action_low), "high:", list(spaces.action_high))
    print("Obs dim:", spaces.obs_dim)

    # Reset
    reset = pb.ResetRequest(render_mode="rgb_array" if args.render else "none", has_seed=(args.seed is not None), seed=(args.seed or 0))
    obs = stub.Reset(reset).observation
    print("Initial obs:", obs)

    # Rollout a few random steps
    rng = np.random.default_rng(0)
    for t in range(args.steps):
        a = rng.uniform(-1.0, 1.0, size=(spaces.action_dim,)).astype(np.float32)
        step = stub.Step(pb.StepRequest(action=list(a)))
        print(f"t={t:03d} reward={step.reward:.4f} term={step.terminated} trunc={step.truncated} obs0={step.observation[:4]}")
        # sleep(1/100.0)
        if args.render and (t % args.render_every == 0):
            frame = stub.Render(pb.RenderRequest(width=args.width, height=args.height))
            with open(f"../../../pic/frame_{t:04d}.png", "wb") as f:
                f.write(frame.rgb_png)
    images_to_video(
        dir_path=args.video_out # args.dir,

    )
    # stub.Close(empty_pb2.Empty())
    channel.close()


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=50051)
    p.add_argument("--steps", type=int, default=500)
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--render", action="store_true")  # 默认渲染
    p.add_argument("--render-every", type=int, default=5)
    p.add_argument("--width", type=int, default=640)
    p.add_argument("--height", type=int, default=480)
    p.add_argument("--video-out", dest="video_out",type=str, default="D:\\workspace\\play_with_pybullet\\pic")
    main(p.parse_args())

